import os
import cv2
import random
import numpy as np
import matplotlib.pyplot as plt

random.seed(33)

# Define dataset paths
dataset_path = "data/wildlife/valid"
image_folder = os.path.join(dataset_path, "images")
label_folder = os.path.join(dataset_path, "labels")

# Class labels dictionary
class_labels = {0: "Buffalo", 1: "Elephant", 2: "Rhino", 3: "Zebra"}
colors = {0: (255, 0, 0), 1: (0, 255, 0), 2: (0, 0, 255), 3: (255, 255, 0)}  # Colors for each class

# Get list of image files
image_files = [f for f in os.listdir(image_folder) if f.endswith(('.jpg', '.png'))]

# Select 6 random images
selected_images = random.sample(image_files, min(9, len(image_files)))

def draw_bboxes(image_path, label_path):
    """Draw bounding boxes on an image using YOLO annotations."""
    # Load image
    image = cv2.imread(image_path)
    height, width, _ = image.shape

    # Read YOLO label file
    with open(label_path, "r") as f:
        lines = f.readlines()

    # Draw bounding boxes
    for line in lines:
        class_id, x_center, y_center, bbox_width, bbox_height = map(float, line.strip().split())

        # Convert YOLO format (normalized) to pixel values
        x_center, y_center = int(x_center * width), int(y_center * height)
        bbox_width, bbox_height = int(bbox_width * width), int(bbox_height * height)

        x_min, y_min = x_center - bbox_width // 2, y_center - bbox_height // 2
        x_max, y_max = x_center + bbox_width // 2, y_center + bbox_height // 2

        # Draw rectangle
        color = colors[int(class_id)]
        cv2.rectangle(image, (x_min, y_min), (x_max, y_max), color, 2)

        # Add class label
        label_text = f"{class_labels[int(class_id)]}"
        cv2.putText(image, label_text, (x_min, y_min - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB for matplotlib

# Plot 6 images with bounding boxes
fig, axes = plt.subplots(3, 3, figsize=(15, 15))

for ax, image_file in zip(axes.flatten(), selected_images):
    image_path = os.path.join(image_folder, image_file)
    label_path = os.path.join(label_folder, image_file.replace(".jpg", ".txt").replace(".png", ".txt"))

    if os.path.exists(label_path):
        processed_image = draw_bboxes(image_path, label_path)
        ax.imshow(processed_image)
        ax.set_title(image_file)
        ax.axis("off")
    else:
        ax.axis("off")

plt.tight_layout()
plt.show()

import os
import cv2
import random
import numpy as np
import matplotlib.pyplot as plt
from ultralytics import YOLO

# Load YOLOv11 pre-trained model
model = YOLO("yolo11s.pt") 

# Function to draw predicted bounding boxes
def draw_predictions(image_path):
    """Runs YOLOv11 inference and draws bounding boxes on the image."""
    # Load image
    image = cv2.imread(image_path)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Run YOLOv11 model on the image
    results = model(image_rgb)

    # Draw bounding boxes
    for result in results:
        boxes = result.boxes.xyxy  # Bounding boxes (x1, y1, x2, y2)
        scores = result.boxes.conf  # Confidence scores
        labels = result.boxes.cls  # Class labels

        for i, box in enumerate(boxes):
            x1, y1, x2, y2 = map(int, box)  # Convert to integers
            label = model.names[int(labels[i])]
            score = scores[i]

            # Draw bounding box
            cv2.rectangle(image_rgb, (x1, y1), (x2, y2), (255, 0, 0), 2)
            cv2.putText(image_rgb, f"{label} {score:.2f}", (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

    return image_rgb  # Return processed image

# Plot 6 images with YOLOv11 predictions
fig, axes = plt.subplots(3, 3, figsize=(15, 15))

for ax, image_file in zip(axes.flatten(), selected_images):
    image_path = os.path.join(image_folder, image_file)
    
    # Process image with YOLOv11
    predicted_image = draw_predictions(image_path)
    
    # Display image
    ax.imshow(predicted_image)
    ax.set_title(f"Predictions for {image_file}")
    ax.axis("off")

plt.tight_layout()
plt.show()

/home/kmcalist/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py:456: UserWarning: Applied workaround for CuDNN issue, install nvrtc.so (Triggered internally at ../aten/src/ATen/native/cudnn/Conv_v8.cpp:80.)
  return F.conv2d(input, weight, bias, self.stride,

0: 448x640 1 elephant, 90.7ms
Speed: 3.8ms preprocess, 90.7ms inference, 57.4ms postprocess per image at shape (1, 3, 448, 640)

0: 416x640 1 zebra, 87.5ms
Speed: 0.9ms preprocess, 87.5ms inference, 1.4ms postprocess per image at shape (1, 3, 416, 640)

0: 480x640 1 cow, 86.2ms
Speed: 2.3ms preprocess, 86.2ms inference, 0.7ms postprocess per image at shape (1, 3, 480, 640)

0: 512x640 2 zebras, 85.7ms
Speed: 2.6ms preprocess, 85.7ms inference, 0.9ms postprocess per image at shape (1, 3, 512, 640)

0: 448x640 1 person, 2 cows, 5.7ms
Speed: 1.1ms preprocess, 5.7ms inference, 0.6ms postprocess per image at shape (1, 3, 448, 640)

0: 448x640 1 cow, 5.5ms
Speed: 1.7ms preprocess, 5.5ms inference, 0.9ms postprocess per image at shape (1, 3, 448, 640)

0: 448x640 1 zebra, 5.2ms
Speed: 1.3ms preprocess, 5.2ms inference, 0.5ms postprocess per image at shape (1, 3, 448, 640)

0: 640x640 4 cows, 2 elephants, 5.5ms
Speed: 0.9ms preprocess, 5.5ms inference, 0.6ms postprocess per image at shape (1, 3, 640, 640)

0: 640x448 1 elephant, 86.3ms
Speed: 0.7ms preprocess, 86.3ms inference, 0.6ms postprocess per image at shape (1, 3, 640, 448)

from ultralytics import YOLO

# Load a pre-trained YOLOv11 model
model = YOLO("yolo11s.pt")  # Use "yolov11s.pt" for better accuracy

# Train the model using your dataset
model.train(
    data="wildlife.yaml",  # Path to your dataset config
    epochs=25,                  # Number of training epochs
    imgsz=640,                  # Image size
    batch=32,                    # Batch size (adjust based on GPU memory)
    workers=12,                   # Number of CPU workers
    device="cuda"                # Use GPU if available, otherwise use "cpu"
)

New https://pypi.org/project/ultralytics/8.3.83 available 😃 Update with 'pip install -U ultralytics'
Ultralytics 8.3.82 🚀 Python-3.10.12 torch-2.1.2+cu118 CUDA:0 (NVIDIA GeForce RTX 3090 Ti, 24245MiB)
engine/trainer: task=detect, mode=train, model=yolo11s.pt, data=wildlife.yaml, epochs=25, time=None, patience=100, batch=32, imgsz=640, save=True, save_period=-1, cache=False, device=cuda, workers=12, project=None, name=train, exist_ok=False, pretrained=True, optimizer=auto, verbose=True, seed=0, deterministic=True, single_cls=False, rect=False, cos_lr=False, close_mosaic=10, resume=False, amp=True, fraction=1.0, profile=False, freeze=None, multi_scale=False, overlap_mask=True, mask_ratio=4, dropout=0.0, val=True, split=val, save_json=False, save_hybrid=False, conf=None, iou=0.7, max_det=300, half=False, dnn=False, plots=True, source=None, vid_stride=1, stream_buffer=False, visualize=False, augment=False, agnostic_nms=False, classes=None, retina_masks=False, embed=None, show=False, save_frames=False, save_txt=False, save_conf=False, save_crop=False, show_labels=True, show_conf=True, show_boxes=True, line_width=None, format=torchscript, keras=False, optimize=False, int8=False, dynamic=False, simplify=True, opset=None, workspace=None, nms=False, lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=7.5, cls=0.5, dfl=1.5, pose=12.0, kobj=1.0, nbs=64, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, bgr=0.0, mosaic=1.0, mixup=0.0, copy_paste=0.0, copy_paste_mode=flip, auto_augment=randaugment, erasing=0.4, crop_fraction=1.0, cfg=None, tracker=botsort.yaml, save_dir=runs/detect/train
Overriding model.yaml nc=80 with nc=4

                   from  n    params  module                                       arguments                     
  0                  -1  1       928  ultralytics.nn.modules.conv.Conv             [3, 32, 3, 2]                 
  1                  -1  1     18560  ultralytics.nn.modules.conv.Conv             [32, 64, 3, 2]                
  2                  -1  1     26080  ultralytics.nn.modules.block.C3k2            [64, 128, 1, False, 0.25]     
  3                  -1  1    147712  ultralytics.nn.modules.conv.Conv             [128, 128, 3, 2]              
  4                  -1  1    103360  ultralytics.nn.modules.block.C3k2            [128, 256, 1, False, 0.25]    
  5                  -1  1    590336  ultralytics.nn.modules.conv.Conv             [256, 256, 3, 2]              
  6                  -1  1    346112  ultralytics.nn.modules.block.C3k2            [256, 256, 1, True]           
  7                  -1  1   1180672  ultralytics.nn.modules.conv.Conv             [256, 512, 3, 2]              
  8                  -1  1   1380352  ultralytics.nn.modules.block.C3k2            [512, 512, 1, True]           
  9                  -1  1    656896  ultralytics.nn.modules.block.SPPF            [512, 512, 5]                 
 10                  -1  1    990976  ultralytics.nn.modules.block.C2PSA           [512, 512, 1]                 
 11                  -1  1         0  torch.nn.modules.upsampling.Upsample         [None, 2, 'nearest']          
 12             [-1, 6]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 13                  -1  1    443776  ultralytics.nn.modules.block.C3k2            [768, 256, 1, False]          
 14                  -1  1         0  torch.nn.modules.upsampling.Upsample         [None, 2, 'nearest']          
 15             [-1, 4]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 16                  -1  1    127680  ultralytics.nn.modules.block.C3k2            [512, 128, 1, False]          
 17                  -1  1    147712  ultralytics.nn.modules.conv.Conv             [128, 128, 3, 2]              
 18            [-1, 13]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 19                  -1  1    345472  ultralytics.nn.modules.block.C3k2            [384, 256, 1, False]          
 20                  -1  1    590336  ultralytics.nn.modules.conv.Conv             [256, 256, 3, 2]              
 21            [-1, 10]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 22                  -1  1   1511424  ultralytics.nn.modules.block.C3k2            [768, 512, 1, True]           
 23        [16, 19, 22]  1    820956  ultralytics.nn.modules.head.Detect           [4, [128, 256, 512]]          
YOLO11s summary: 181 layers, 9,429,340 parameters, 9,429,324 gradients, 21.6 GFLOPs

Transferred 493/499 items from pretrained weights
WARNING ⚠️ Comet installed but not initialized correctly, not logging this run. start() got an unexpected keyword argument 'project_name'
WARNING ⚠️ NeptuneAI installed but not initialized correctly, not logging this run. 

----NeptuneMissingApiTokenException-------------------------------------------

The Neptune client couldn't find your API token.

You can get it here:
    - https://app.neptune.ai/get_my_api_token

There are two options to add it:
    - specify it in your code
    - set an environment variable in your operating system.

CODE
Pass the token to the init_run() function via the api_token argument:
    neptune.init_run(project='WORKSPACE_NAME/PROJECT_NAME', api_token='YOUR_API_TOKEN')

ENVIRONMENT VARIABLE (Recommended option)
or export or set an environment variable depending on your operating system:

    Linux/Unix
    In your terminal run:
        export NEPTUNE_API_TOKEN="YOUR_API_TOKEN"

    Windows
    In your CMD run:
        set NEPTUNE_API_TOKEN="YOUR_API_TOKEN"

and skip the api_token argument of the init_run() function:
    neptune.init_run(project='WORKSPACE_NAME/PROJECT_NAME')

You may also want to check the following docs pages:
    - https://docs.neptune.ai/setup/setting_api_token/

Need help?-> https://docs.neptune.ai/getting_help

TensorBoard: Start with 'tensorboard --logdir runs/detect/train', view at http://localhost:6006/
Freezing layer 'model.23.dfl.conv.weight'
AMP: running Automatic Mixed Precision (AMP) checks...

[neptune] [warning] NeptuneWarning: The following monitoring options are disabled by default in interactive sessions: 'capture_stdout', 'capture_stderr', 'capture_traceback', and 'capture_hardware_metrics'. To enable them, set each parameter to 'True' when initializing the run. The monitoring will continue until you call run.stop() or the kernel stops. Also note: Your source files can only be tracked if you pass the path(s) to the 'source_code' argument. For help, see the Neptune docs: https://docs.neptune.ai/logging/source_code/

AMP: checks passed ✅

train: Scanning /home/kmcalist/QTM447/Spring2025/Lectures/Lecture15/data/wildlife/train/labels.cache... 1052 images, 0 backgrounds, 0 corrupt: 100%|██████████| 1052/1052 [00:00<?, ?it/s]
val: Scanning /home/kmcalist/QTM447/Spring2025/Lectures/Lecture15/data/wildlife/valid/labels.cache... 225 images, 0 backgrounds, 0 corrupt: 100%|██████████| 225/225 [00:00<?, ?it/s]

Plotting labels to runs/detect/train/labels.jpg... 
optimizer: 'optimizer=auto' found, ignoring 'lr0=0.01' and 'momentum=0.937' and determining best 'optimizer', 'lr0' and 'momentum' automatically... 
optimizer: AdamW(lr=0.00125, momentum=0.9) with parameter groups 81 weight(decay=0.0), 88 weight(decay=0.0005), 87 bias(decay=0.0)
TensorBoard: model graph visualization added ✅
Image sizes 640 train, 640 val
Using 12 dataloader workers
Logging results to runs/detect/train
Starting training for 25 epochs...

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size

       1/25      8.42G     0.8596      2.283      1.256         98        640: 100%|██████████| 33/33 [00:06<00:00,  5.11it/s]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 4/4 [00:01<00:00,  3.56it/s]

                   all        225        379      0.604      0.604      0.691      0.472

      Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size

       2/25      8.18G      1.008      1.331      1.328        115        640: 100%|██████████| 33/33 [00:05<00:00,  6.07it/s]
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 4/4 [00:00<00:00,  6.54it/s]

better_model = YOLO("runs/detect/train/weights/best.pt")

import os
import cv2
import random
import numpy as np
import matplotlib.pyplot as plt
from ultralytics import YOLO

# Function to draw predicted bounding boxes
def draw_predictions(image_path):
    """Runs YOLOv11 inference and draws bounding boxes on the image."""
    # Load image
    image = cv2.imread(image_path)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Run YOLOv11 model on the image
    results = better_model(image_rgb)

    # Draw bounding boxes
    for result in results:
        boxes = result.boxes.xyxy  # Bounding boxes (x1, y1, x2, y2)
        scores = result.boxes.conf  # Confidence scores
        labels = result.boxes.cls  # Class labels

        for i, box in enumerate(boxes):
            x1, y1, x2, y2 = map(int, box)  # Convert to integers
            label = model.names[int(labels[i])]
            score = scores[i]

            # Draw bounding box
            cv2.rectangle(image_rgb, (x1, y1), (x2, y2), (255, 0, 0), 2)
            cv2.putText(image_rgb, f"{label} {score:.2f}", (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

    return image_rgb  # Return processed image

# Plot 6 images with YOLOv11 predictions
fig, axes = plt.subplots(3, 3, figsize=(15, 15))

for ax, image_file in zip(axes.flatten(), selected_images):
    image_path = os.path.join(image_folder, image_file)
    
    # Process image with YOLOv11
    predicted_image = draw_predictions(image_path)
    
    # Display image
    ax.imshow(predicted_image)
    ax.set_title(f"Predictions for {image_file}")
    ax.axis("off")

plt.tight_layout()
plt.show()

0: 448x640 2 elephants, 8.4ms
Speed: 0.6ms preprocess, 8.4ms inference, 0.7ms postprocess per image at shape (1, 3, 448, 640)

0: 416x640 1 zebra, 8.7ms
Speed: 0.8ms preprocess, 8.7ms inference, 0.6ms postprocess per image at shape (1, 3, 416, 640)

0: 480x640 1 rhino, 8.5ms
Speed: 1.0ms preprocess, 8.5ms inference, 0.6ms postprocess per image at shape (1, 3, 480, 640)

0: 512x640 2 zebras, 8.5ms
Speed: 0.9ms preprocess, 8.5ms inference, 0.6ms postprocess per image at shape (1, 3, 512, 640)

0: 448x640 1 rhino, 7.5ms
Speed: 0.8ms preprocess, 7.5ms inference, 0.6ms postprocess per image at shape (1, 3, 448, 640)

0: 448x640 1 buffalo, 7.1ms
Speed: 0.9ms preprocess, 7.1ms inference, 0.6ms postprocess per image at shape (1, 3, 448, 640)

0: 448x640 1 zebra, 7.1ms
Speed: 0.6ms preprocess, 7.1ms inference, 0.6ms postprocess per image at shape (1, 3, 448, 640)

0: 640x640 6 rhinos, 8.8ms
Speed: 1.0ms preprocess, 8.8ms inference, 0.8ms postprocess per image at shape (1, 3, 640, 640)

0: 640x448 1 elephant, 8.9ms
Speed: 0.6ms preprocess, 8.9ms inference, 0.7ms postprocess per image at shape (1, 3, 640, 448)

import pandas as pd
import numpy as np
import os

#Organize oxford pets into folder structures

np.random.seed(42)

# Get list of image files
image_folder = "data/oxford-iiit-pet/images"
image_files = [f for f in os.listdir(image_folder) if f.endswith(('.jpg', '.png'))]

# Create dataframe
df = pd.DataFrame(image_files, columns=['filename'])

# Add label column
df['label'] = df['filename'].apply(lambda x: 'cat' if x[0].isupper() else 'dog')

# Add set column
df['set'] = np.random.choice([1, 0], size=len(df), p=[0.5, 0.5])

print(df.head())

                           filename label  set
0        german_shorthaired_184.jpg   dog    1
1                    Birman_120.jpg   cat    0
2             great_pyrenees_20.jpg   dog    0
3                     samoyed_6.jpg   dog    0
4  american_pit_bull_terrier_28.jpg   dog    1

import os

# Define the base path
base_path = "data/oxford_pets"

# Define the folder structure
folders = [
    "train/cats", "train/dogs",
    "valid/cats", "valid/dogs",
    "train_trimaps/cats", "train_trimaps/dogs",
    "valid_trimaps/cats", "valid_trimaps/dogs"
]

# Create the folders
for folder in folders:
    os.makedirs(os.path.join(base_path, folder), exist_ok=True)

print("Folder structure created successfully.")

Folder structure created successfully.

import shutil

# Define the trimap folder
trimap_folder = "data/oxford-iiit-pet/annotations/trimaps"

# Function to copy files to the appropriate folder
def copy_files(row):
    image_file = row['filename']
    label = row['label']
    set_type = 'train' if row['set'] == 1 else 'valid'
    
    # Define source and destination paths
    image_src = os.path.join(image_folder, image_file)
    trimap_src = os.path.join(trimap_folder, image_file.replace('.jpg', '.png'))
    
    if os.path.exists(trimap_src):
        image_dst = os.path.join(base_path, f"{set_type}/{label}s", image_file)
        trimap_dst = os.path.join(base_path, f"{set_type}_trimaps/{label}s", image_file.replace('.jpg', '.png'))
        
        # Copy image and trimap
        shutil.copy(image_src, image_dst)
        shutil.copy(trimap_src, trimap_dst)

# Apply the function to each row in the dataframe
df.apply(copy_files, axis=1)

print("Files copied successfully.")

Files copied successfully.

import os
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image
import numpy as np

# Define class mapping
CLASS_MAP = {"cats": 0, "dogs": 1}  # 0 for cats, 1 for dogs

# Define a consistent image size
IMAGE_SIZE = (256, 256)

class OxfordPetsDataset(Dataset):
    def __init__(self, image_dir, trimap_dir, transform=None):
        self.image_dir = image_dir
        self.trimap_dir = trimap_dir
        self.transform = transform

        # List all image files
        self.image_files = sorted(os.listdir(image_dir))

        # Define transformations to ensure consistent size for trimaps (without normalization)
        self.trimap_transform = transforms.Compose([
            transforms.Resize(IMAGE_SIZE, interpolation=Image.NEAREST),  # Resize while preserving labels
        ])

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        # Get filename
        image_filename = self.image_files[idx]
        image_path = os.path.join(self.image_dir, image_filename)
        trimap_path = os.path.join(self.trimap_dir, image_filename.replace(".jpg", ".png"))

        # Load image and trimap
        image = Image.open(image_path).convert("RGB")
        trimap = Image.open(trimap_path).convert("L")  # Trimap is grayscale

        # Resize trimap to (256, 256)
        trimap = self.trimap_transform(trimap)
        trimap = np.array(trimap, dtype=np.uint8)  # Ensure integer encoding

        # Get class label from folder name (cats or dogs)
        class_name = os.path.basename(os.path.dirname(image_path))
        class_label = CLASS_MAP[class_name]  # 0 for cats, 1 for dogs

        # Modify trimap coding (3-class)
        new_trimap = np.zeros_like(trimap, dtype=np.uint8)  # Initialize new trimap
        new_trimap[trimap == 1] = 1  # Outline (both cats & dogs)
        new_trimap[trimap == 2] = 2  # Object (both cats & dogs)

        # Convert trimap to LongTensor before returning
        new_trimap = torch.tensor(new_trimap, dtype=torch.long)  # Ensure it's stored correctly

        # Apply transformations to image (but NOT trimap)
        if self.transform:
            image = self.transform(image)

        return image, new_trimap, torch.tensor(class_label, dtype=torch.long)

import os
import torch
from torch.utils.data import Dataset, DataLoader, ConcatDataset
from torchvision import transforms

# Define dataset paths for images and trimaps
train_cats_images = "data/oxford_pets/train/cats"
train_dogs_images = "data/oxford_pets/train/dogs"
train_cats_trimaps = "data/oxford_pets/train_trimaps/cats"
train_dogs_trimaps = "data/oxford_pets/train_trimaps/dogs"

valid_cats_images = "data/oxford_pets/valid/cats"
valid_dogs_images = "data/oxford_pets/valid/dogs"
valid_cats_trimaps = "data/oxford_pets/valid_trimaps/cats"
valid_dogs_trimaps = "data/oxford_pets/valid_trimaps/dogs"

# Define transformations (apply only to images)
data_transforms = transforms.Compose([
    transforms.Resize((256, 256)),  # Ensure consistent image size
    transforms.ToTensor(),  # Convert image to PyTorch tensor (normalizes to [0,1])
])

# Create Dataset Instances
train_cats_dataset = OxfordPetsDataset(train_cats_images, train_cats_trimaps, transform=data_transforms)
train_dogs_dataset = OxfordPetsDataset(train_dogs_images, train_dogs_trimaps, transform=data_transforms)

valid_cats_dataset = OxfordPetsDataset(valid_cats_images, valid_cats_trimaps, transform=data_transforms)
valid_dogs_dataset = OxfordPetsDataset(valid_dogs_images, valid_dogs_trimaps, transform=data_transforms)

# Merge datasets
from torch.utils.data import ConcatDataset
train_dataset = ConcatDataset([train_cats_dataset, train_dogs_dataset])
valid_dataset = ConcatDataset([valid_cats_dataset, valid_dogs_dataset])

# Create DataLoaders
batch_size = 256
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False)

print(f"Train set: {len(train_loader.dataset)} samples")
print(f"Valid set: {len(valid_loader.dataset)} samples")

Train set: 3723 samples
Valid set: 3667 samples

import matplotlib.pyplot as plt
import random
import numpy as np

# Select a random index from the validation dataset
random_idx = random.randint(0, len(valid_dataset) - 1)

# Get the image, trimap, and class label
image, trimap, class_label = valid_dataset[random_idx]

# Convert image tensor to NumPy format for visualization
img_np = image.permute(1, 2, 0).numpy()  # Convert from (C, H, W) to (H, W, C)

# Convert trimap tensor to NumPy
trimap_np = trimap.cpu().numpy()

np.unique(trimap_np)

array([0, 1, 2])

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.colors as mcolors

# Define class-to-color mapping for trimap
trimap_colors = {
    0: (0, 0, 0),       # Background - Black
    1: (0, 0, 1),       # Dog Outline - Blue
    2: (0, 1, 0),       # Dog Object - Green
    3: (1, 0, 0),       # Cat Outline - Red
    4: (1, 1, 0),       # Cat Object - Yellow
}

# Create a ListedColormap for visualization
cmap = mcolors.ListedColormap([trimap_colors[i] for i in range(len(trimap_colors))])

# Get a batch of training data
images, trimaps, labels = next(iter(train_loader))

# Show the first 4 images and trimaps
fig, axes = plt.subplots(2, 4, figsize=(12, 6))

for i in range(4):
    # Convert tensor image to NumPy format
    img_np = images[i].permute(1, 2, 0).cpu().numpy()
    img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min())  # Normalize for display

    # Convert trimap tensor to NumPy
    trimap_np = trimaps[i].cpu().numpy()

    # Image display
    axes[0, i].imshow(img_np)
    axes[0, i].set_title(f"Class: {'Dog' if labels[i] == 1 else 'Cat'}")
    axes[0, i].axis("off")

    # Trimap display with custom colors
    axes[1, i].imshow(trimap_np, cmap=cmap, vmin=0, vmax=len(trimap_colors)-1)
    axes[1, i].set_title("Trimap (Labeled)")
    axes[1, i].axis("off")

# Create legend for the trimap colors
fig.subplots_adjust(right=0.85)
cbar_ax = fig.add_axes([0.88, 0.15, 0.02, 0.7])
cbar = plt.colorbar(mappable=plt.cm.ScalarMappable(cmap=cmap), cax=cbar_ax, ticks=np.arange(len(trimap_colors)))
cbar.ax.set_yticklabels(["Background", "Dog Outline", "Dog", "Cat Outline", "Cat"])  # Custom labels

plt.tight_layout()
plt.show()

/tmp/ipykernel_380213/3389472999.py:47: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
  plt.tight_layout()

import torch
import torch.nn as nn
import torchvision.models as models

class UNetMobileNet(nn.Module):
    def __init__(self, num_classes=5, pretrained=True, freeze_encoder=True, use_skip_connections=True):
        super(UNetMobileNet, self).__init__()

        self.use_skip_connections = use_skip_connections  # Toggle skip connections

        # Load MobileNetV2 as encoder
        mobilenet = models.mobilenet_v2(pretrained=pretrained)
        encoder_layers = list(mobilenet.features.children())

        # Encoder blocks
        self.encoder1 = nn.Sequential(*encoder_layers[:2])   # 128 × 128
        self.encoder2 = nn.Sequential(*encoder_layers[2:4])  # 64 × 64
        self.encoder3 = nn.Sequential(*encoder_layers[4:7])  # 32 × 32
        self.encoder4 = nn.Sequential(*encoder_layers[7:14]) # 16 × 16
        self.encoder5 = nn.Sequential(*encoder_layers[14:])  # 8 × 8

        # Freeze encoder if specified
        if freeze_encoder:
            for param in self.encoder1.parameters():
                param.requires_grad = False
            for param in self.encoder2.parameters():
                param.requires_grad = False
            for param in self.encoder3.parameters():
                param.requires_grad = False
            for param in self.encoder4.parameters():
                param.requires_grad = False
            for param in self.encoder5.parameters():
                param.requires_grad = False

        # Bottleneck layer
        self.bottleneck = nn.Conv2d(1280, 512, kernel_size=1)

        # Decoder (Transposed Convolutions for upsampling)
        self.decoder4 = self._upsample(512, 96)  # 16 × 16
        self.decoder3 = self._upsample(96, 32)   # 32 × 32
        self.decoder2 = self._upsample(32, 24)   # 64 × 64
        self.decoder1 = self._upsample(24, 16)   # 128 × 128

        # **Extra upsampling to reach 256 × 256**
        self.final_up = self._upsample(16, 16)   # 256 × 256

        self.final_conv = nn.Conv2d(16, num_classes, kernel_size=1)  # Final segmentation layer

    def _upsample(self, in_channels, out_channels):
        """Helper function to create an upsampling block using transposed convolutions."""
        return nn.Sequential(
            nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        # Encoder
        e1 = self.encoder1(x)  # 128 × 128
        e2 = self.encoder2(e1) # 64 × 64
        e3 = self.encoder3(e2) # 32 × 32
        e4 = self.encoder4(e3) # 16 × 16
        e5 = self.encoder5(e4) # 8 × 8

        # Bottleneck layer
        b = self.bottleneck(e5)

        # Decoder with optional skip connections
        d4 = self.decoder4(b) + e4 if self.use_skip_connections else self.decoder4(b)  # 16 × 16
        d3 = self.decoder3(d4) + e3 if self.use_skip_connections else self.decoder3(d4)  # 32 × 32
        d2 = self.decoder2(d3) + e2 if self.use_skip_connections else self.decoder2(d3)  # 64 × 64
        d1 = self.decoder1(d2) + e1 if self.use_skip_connections else self.decoder1(d2)  # 128 × 128

        d0 = self.final_up(d1)  # **Extra upsampling to reach 256 × 256**

        # Final segmentation map
        return self.final_conv(d0)  # Shape: (batch_size, num_classes, 256, 256)

import networkx as nx
import matplotlib.pyplot as plt

def visualize_unet_with_correct_skips(use_skip_connections=True):
    """
    Creates a U-Net architecture visualization with input/output sizes and channels at each layer.
    
    Args:
        use_skip_connections (bool): Whether to show skip connections.
    
    Returns:
        fig: Matplotlib figure to be displayed inline in Jupyter Notebook.
    """
    G = nx.DiGraph()

    # Define encoder layers with (Height x Width, Channels)
    encoder_layers = {
        "Input (256x256, 3)": (0, 5),
        "Enc1 (128x128, 16)": (1, 4),
        "Enc2 (64x64, 24)": (2, 3),
        "Enc3 (32x32, 32)": (3, 2),
        "Enc4 (16x16, 96)": (4, 1),
        "Bottleneck (8x8, 512)": (5, 0)
    }

    # Define decoder layers with corresponding output sizes and channels
    decoder_layers = {
        "Dec4 (16x16, 96)": (6, 1),
        "Dec3 (32x32, 32)": (7, 2),
        "Dec2 (64x64, 24)": (8, 3),
        "Dec1 (128x128, 16)": (9, 4),
        "Final Up (256x256, 16)": (10, 5),
        "Output (256x256, 5)": (11, 5.5)
    }

    # Add encoder and decoder nodes
    for layer, pos in {**encoder_layers, **decoder_layers}.items():
        G.add_node(layer, pos=pos)

    # Connect encoder layers sequentially
    encoder_keys = list(encoder_layers.keys())
    for i in range(len(encoder_keys) - 1):
        G.add_edge(encoder_keys[i], encoder_keys[i + 1])

    # Connect bottleneck to decoder
    G.add_edge("Bottleneck (8x8, 512)", "Dec4 (16x16, 96)")

    # Connect decoder layers sequentially
    decoder_keys = list(decoder_layers.keys())
    for i in range(len(decoder_keys) - 2):
        G.add_edge(decoder_keys[i], decoder_keys[i + 1])

    # Connect final output
    G.add_edge("Final Up (256x256, 16)", "Output (256x256, 5)")

    # Add **correct** skip connections if enabled
    if use_skip_connections:
        skip_connections = {
            "Enc1 (128x128, 16)": "Dec1 (128x128, 16)",
            "Enc2 (64x64, 24)": "Dec2 (64x64, 24)",
            "Enc3 (32x32, 32)": "Dec3 (32x32, 32)",
            "Enc4 (16x16, 96)": "Dec4 (16x16, 96)"
        }
        for enc, dec in skip_connections.items():
            G.add_edge(enc, dec, color="red", style="dashed")

    # Extract positions for visualization
    pos = nx.get_node_attributes(G, "pos")

    # Create a Matplotlib figure
    fig, ax = plt.subplots(figsize=(12, 6))
    edges = G.edges()

    # Color edges differently for skip connections
    edge_colors = ["red" if G[u][v].get("color") == "red" else "black" for u, v in edges]
    edge_styles = ["dashed" if G[u][v].get("style") == "dashed" else "solid" for u, v in edges]

    # Draw nodes
    nx.draw(G, pos, with_labels=True, node_color="lightblue", node_size=2500, font_size=8, font_weight="bold", ax=ax)

    # Draw normal and skip edges separately to allow different styles
    for i, (u, v) in enumerate(edges):
        nx.draw_networkx_edges(G, pos, edgelist=[(u, v)], edge_color=edge_colors[i], style=edge_styles[i], width=2, ax=ax)

    ax.set_title(f"U-Net Architecture {'(With Skip Connections)' if use_skip_connections else '(Without Skip Connections)'}")
    
    return fig

# Display the images inline in Jupyter Notebook
fig1 = visualize_unet_with_correct_skips(use_skip_connections=True)  # With Skip Connections
fig2 = visualize_unet_with_correct_skips(use_skip_connections=False) # Without Skip Connections

# Show figures in Jupyter Notebook
plt.show()

import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.optim as optim
import torchmetrics

class UNetLightning(pl.LightningModule):
    def __init__(self, num_classes=5, lr=1e-3, freeze_encoder=True, use_skip_connections=True):
        super(UNetLightning, self).__init__()

        # Load U-Net with correct output shape (256x256)
        self.model = UNetMobileNet(
            num_classes=num_classes, 
            freeze_encoder=freeze_encoder, 
            use_skip_connections=use_skip_connections
        )

        # Loss function (Cross-entropy for multi-class segmentation)
        self.criterion = nn.CrossEntropyLoss()

        # IoU (Intersection over Union) metric
        self.iou = torchmetrics.JaccardIndex(task="multiclass", num_classes=num_classes)

        # Learning rate
        self.lr = lr

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        """ Training step: Computes loss and IoU for segmentation. """
        images, trimaps, _ = batch  # Extract inputs and ground truth masks
        logits = self.model(images)  # Forward pass
        loss = self.criterion(logits, trimaps.squeeze(1).long())  # Squeeze trimap to shape (batch_size, height, width)

        # Compute IoU (Jaccard Index)
        iou = self.iou(torch.argmax(logits, dim=1), trimaps.squeeze(1).long())

        # Log loss and IoU for training
        self.log("train_loss", loss, prog_bar=True)
        self.log("train_iou", iou, prog_bar=True)

        return loss

    def validation_step(self, batch, batch_idx):
        """ Validation step: Evaluates loss and IoU on the validation set. """
        images, trimaps, _ = batch
        logits = self.model(images)
        loss = self.criterion(logits, trimaps.squeeze(1).long())  # Squeeze trimap to shape (batch_size, height, width)

        # Compute IoU
        iou = self.iou(torch.argmax(logits, dim=1), trimaps.squeeze(1).long())

        # Log metrics for validation
        self.log("val_loss", loss, prog_bar=True)
        self.log("val_iou", iou, prog_bar=True)

        return loss


    def predict_step(self, batch, batch_idx):
        """ Prediction step: Runs inference on new images. """
        images, _, _ = batch  # We only need images for prediction
        logits = self.model(images)  # Forward pass
        preds = torch.argmax(logits, dim=1)  # Get predicted class labels

        return preds  # Return predicted segmentation masks

    def configure_optimizers(self):
        """ Optimizer and Learning Rate Scheduler """
        optimizer = optim.Adam(self.parameters(), lr=self.lr)
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=3)

        return {"optimizer": optimizer, "lr_scheduler": scheduler, "monitor": "val_loss"}

from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.loggers import CSVLogger
# Set up logging and early stopping
csv_logger = CSVLogger(save_dir='logs/', name='UNetNoSkips', version="")
early_stop_callback = EarlyStopping(monitor='val_loss', patience=25, verbose=True, mode="min")

# Create the model instance
model = UNetLightning(use_skip_connections = False)

# Assume train_loader and val_loader are defined DataLoaders
trainer = pl.Trainer(
    max_epochs=10,
    logger=csv_logger,
    callbacks=[early_stop_callback]
)

trainer.fit(model, train_loader, valid_loader)

# Save the final model state
trainer.save_checkpoint('logs/UNetNoSkips/final_model.ckpt')

/home/kmcalist/.local/lib/python3.10/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.
  warnings.warn(
/home/kmcalist/.local/lib/python3.10/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=MobileNet_V2_Weights.IMAGENET1K_V1`. You can also use `weights=MobileNet_V2_Weights.DEFAULT` to get the most up-to-date weights.
  warnings.warn(msg)
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
/home/kmcalist/.local/lib/python3.10/site-packages/lightning_fabric/loggers/csv_logs.py:269: Experiment logs directory logs/UNetNoSkips/ exists and is not empty. Previous log files in this directory will be deleted when the new ones are saved!
/home/kmcalist/.local/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:652: Checkpoint directory logs/UNetNoSkips/checkpoints exists and is not empty.
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

  | Name      | Type                   | Params | Mode 
-------------------------------------------------------------
0 | model     | UNetMobileNet          | 3.8 M  | train
1 | criterion | CrossEntropyLoss       | 0      | train
2 | iou       | MulticlassJaccardIndex | 0      | train
-------------------------------------------------------------
1.6 M     Trainable params
2.2 M     Non-trainable params
3.8 M     Total params
15.361    Total estimated model params size (MB)

Sanity Checking: |          | 0/? [00:00<?, ?it/s]

/home/kmcalist/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:424: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=23` in the `DataLoader` to improve performance.
/home/kmcalist/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:424: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=23` in the `DataLoader` to improve performance.
/home/kmcalist/.local/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py:298: The number of training batches (15) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.

Training: |          | 0/? [00:00<?, ?it/s]

Validation: |          | 0/? [00:00<?, ?it/s]

Metric val_loss improved. New best score: 1.304

Validation: |          | 0/? [00:00<?, ?it/s]

Metric val_loss improved by 0.245 >= min_delta = 0.0. New best score: 1.059

Validation: |          | 0/? [00:00<?, ?it/s]

Metric val_loss improved by 0.172 >= min_delta = 0.0. New best score: 0.887

from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.loggers import CSVLogger
# Set up logging and early stopping
csv_logger = CSVLogger(save_dir='logs/', name='UNetSkips', version="")
early_stop_callback = EarlyStopping(monitor='val_loss', patience=25, verbose=True, mode="min")

# Create the model instance
model = UNetLightning(use_skip_connections = True)

# Assume train_loader and val_loader are defined DataLoaders
trainer = pl.Trainer(
    max_epochs=5,
    logger=csv_logger,
    callbacks=[early_stop_callback]
)

trainer.fit(model, train_loader, valid_loader)

# Save the final model state
trainer.save_checkpoint('logs/UNetSkips/final_model.ckpt')

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
/home/kmcalist/.local/lib/python3.10/site-packages/lightning_fabric/loggers/csv_logs.py:269: Experiment logs directory logs/UNetSkips/ exists and is not empty. Previous log files in this directory will be deleted when the new ones are saved!
/home/kmcalist/.local/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:652: Checkpoint directory logs/UNetSkips/checkpoints exists and is not empty.
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

  | Name      | Type                   | Params | Mode 
-------------------------------------------------------------
0 | model     | UNetMobileNet          | 3.8 M  | train
1 | criterion | CrossEntropyLoss       | 0      | train
2 | iou       | MulticlassJaccardIndex | 0      | train
-------------------------------------------------------------
1.6 M     Trainable params
2.2 M     Non-trainable params
3.8 M     Total params
15.361    Total estimated model params size (MB)

Sanity Checking: |          | 0/? [00:00<?, ?it/s]

Training: |          | 0/? [00:00<?, ?it/s]

Validation: |          | 0/? [00:00<?, ?it/s]

Metric val_loss improved. New best score: 1.336

Validation: |          | 0/? [00:00<?, ?it/s]

Metric val_loss improved by 0.289 >= min_delta = 0.0. New best score: 1.047

Validation: |          | 0/? [00:00<?, ?it/s]

Metric val_loss improved by 0.035 >= min_delta = 0.0. New best score: 1.012

Validation: |          | 0/? [00:00<?, ?it/s]

import torch
import pytorch_lightning as pl

# Load trained models
model_with_skips = UNetLightning.load_from_checkpoint(
    'logs/UNetSkips/final_model.ckpt',
    use_skip_connections=True  # Ensure it matches training
)

model_without_skips = UNetLightning.load_from_checkpoint(
    'logs/UNetNoSkips/final_model.ckpt',
    use_skip_connections=False  # Ensure it matches training
)

# Set models to evaluation mode
model_with_skips.eval()
model_without_skips.eval()

print("done")

done

/home/kmcalist/.local/lib/python3.10/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.
  warnings.warn(
/home/kmcalist/.local/lib/python3.10/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=MobileNet_V2_Weights.IMAGENET1K_V1`. You can also use `weights=MobileNet_V2_Weights.DEFAULT` to get the most up-to-date weights.
  warnings.warn(msg)

import random

# Get the validation dataset from the DataLoader
valid_dataset = valid_loader.dataset

# Randomly select 10 unique indices
random_indices = random.sample(range(len(valid_dataset)), 10)

# Extract the corresponding images (without labels)
random_samples = [valid_dataset[i] for i in random_indices]

# Convert images to batch format for model prediction
images = torch.stack([sample[0] for sample in random_samples])  # Image tensors
trimaps = torch.stack([sample[1] for sample in random_samples])  # Ground truth trimaps
classes = torch.stack([sample[2] for sample in random_samples])

# Create DataLoader for prediction
pred_loader = torch.utils.data.DataLoader(list(zip(images, trimaps, classes)), batch_size=10)

# Define Lightning trainer (No training, just prediction mode)
trainer = pl.Trainer(accelerator="gpu" if torch.cuda.is_available() else "cpu")

# Run predictions with trainer.predict()
preds_with_skips = trainer.predict(model_with_skips, dataloaders=pred_loader)
preds_without_skips = trainer.predict(model_without_skips, dataloaders=pred_loader)

# Convert list of batch tensors to a single tensor
preds_with_skips = torch.cat(preds_with_skips, dim=0)
preds_without_skips = torch.cat(preds_without_skips, dim=0)

# Convert logits to class labels
preds_with_skips = torch.argmax(preds_with_skips, dim=1).cpu().numpy()
preds_without_skips = torch.argmax(preds_without_skips, dim=1).cpu().numpy()
trimaps = trimaps.squeeze(1).cpu().numpy()

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
2025-03-04 14:16:52.935266: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2025-03-04 14:16:52.935290: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2025-03-04 14:16:52.936145: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2025-03-04 14:16:52.940029: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2025-03-04 14:16:53.520495: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
/home/kmcalist/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:424: The 'predict_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=23` in the `DataLoader` to improve performance.

Predicting: |          | 0/? [00:00<?, ?it/s]

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Predicting: |          | 0/? [00:00<?, ?it/s]

preds_without_skips

array([[0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       ...,
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0]])

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.colors as mcolors

# Define class-to-color mapping for trimap
trimap_colors = {
    0: (0, 0, 0),   # Background - Black
    1: (1, 0, 0),   # Outline - Red
    2: (0, 1, 0),   # Object - Green
}

# Create a ListedColormap for visualization
cmap = mcolors.ListedColormap([trimap_colors[i] for i in range(len(trimap_colors))])

# Show the first 10 images, ground truth trimaps, and predicted trimaps
fig, axes = plt.subplots(10, 4, figsize=(12, 30))
axes[0, 0].set_title("Input Image")
axes[0, 1].set_title("Ground Truth Trimap")
axes[0, 2].set_title("Prediction (With Skips)")
axes[0, 3].set_title("Prediction (No Skips)")

for i in range(10):
    # Convert image tensor to NumPy format
    img_np = images[i].permute(1, 2, 0).cpu().numpy()
    img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min())  # Normalize for display

    # Convert trimap tensors to NumPy
    trimap_np = trimaps[i]
    pred_skips_np = preds_with_skips[i]
    pred_no_skips_np = preds_without_skips[i]

    axes[i, 0].imshow(img_np)
    axes[i, 0].axis("off")

    axes[i, 1].imshow(trimap_np, cmap=cmap, vmin=0, vmax=len(trimap_colors)-1)
    axes[i, 1].axis("off")

    axes[i, 2].imshow(pred_skips_np, cmap=cmap, vmin=0, vmax=len(trimap_colors)-1)
    axes[i, 2].axis("off")

    axes[i, 3].imshow(pred_no_skips_np, cmap=cmap, vmin=0, vmax=len(trimap_colors)-1)
    axes[i, 3].axis("off")

plt.tight_layout()
plt.show()

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[56], line 38
     35 axes[i, 1].imshow(trimap_np, cmap=cmap, vmin=0, vmax=len(trimap_colors)-1)
     36 axes[i, 1].axis("off")
---> 38 axes[i, 2].imshow(pred_skips_np, cmap=cmap, vmin=0, vmax=len(trimap_colors)-1)
     39 axes[i, 2].axis("off")
     41 axes[i, 3].imshow(pred_no_skips_np, cmap=cmap, vmin=0, vmax=len(trimap_colors)-1)

File ~/.local/lib/python3.10/site-packages/matplotlib/__init__.py:1521, in _preprocess_data.<locals>.inner(ax, data, *args, **kwargs)
   1518 @functools.wraps(func)
   1519 def inner(ax, *args, data=None, **kwargs):
   1520     if data is None:
-> 1521         return func(
   1522             ax,
   1523             *map(cbook.sanitize_sequence, args),
   1524             **{k: cbook.sanitize_sequence(v) for k, v in kwargs.items()})
   1526     bound = new_sig.bind(ax, *args, **kwargs)
   1527     auto_label = (bound.arguments.get(label_namer)
   1528                   or bound.kwargs.get(label_namer))

File ~/.local/lib/python3.10/site-packages/matplotlib/axes/_axes.py:5945, in Axes.imshow(self, X, cmap, norm, aspect, interpolation, alpha, vmin, vmax, colorizer, origin, extent, interpolation_stage, filternorm, filterrad, resample, url, **kwargs)
   5942 if aspect is not None:
   5943     self.set_aspect(aspect)
-> 5945 im.set_data(X)
   5946 im.set_alpha(alpha)
   5947 if im.get_clip_path() is None:
   5948     # image does not already have clipping set, clip to Axes patch

File ~/.local/lib/python3.10/site-packages/matplotlib/image.py:675, in _ImageBase.set_data(self, A)
    673 if isinstance(A, PIL.Image.Image):
    674     A = pil_to_array(A)  # Needed e.g. to apply png palette.
--> 675 self._A = self._normalize_image_array(A)
    676 self._imcache = None
    677 self.stale = True

File ~/.local/lib/python3.10/site-packages/matplotlib/image.py:643, in _ImageBase._normalize_image_array(A)
    641     A = A.squeeze(-1)  # If just (M, N, 1), assume scalar and apply colormap.
    642 if not (A.ndim == 2 or A.ndim == 3 and A.shape[-1] in [3, 4]):
--> 643     raise TypeError(f"Invalid shape {A.shape} for image data")
    644 if A.ndim == 3:
    645     # If the input data has values outside the valid range (after
    646     # normalisation), we issue a warning and then clip X to the bounds
    647     # - otherwise casting wraps extreme values, hiding outliers and
    648     # making reliable interpretation impossible.
    649     high = 255 if np.issubdtype(A.dtype, np.integer) else 1

TypeError: Invalid shape (256,) for image data

YOLO for African Wildlife Object Detection¶

Fine Tuning¶

Semantic Segmentation with a Custom U-Net¶